package util

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/hex"
	"encoding/pem"
	"github.com/pkg/errors"
	"math/big"
	"os"
	"regexp"
)

const (
	numeric = "0123456789"
	upper   = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
	lower   = "abcdefghijklmnopqrstuvwxyz"
)

var (
	regexPassword = regexp.MustCompile("^[\x20-\x7E]*$")
)

// RandomAlphanumericString generates a random password of some fixed length.
func RandomAlphanumericString(strLen int) string {
	result := make([]byte, strLen)
	typeList := randomTypeList(strLen)
	for i, t := range typeList {
		result[i] = randomValue(t)
	}
	return string(result)
}

func randomTypeList(length int) []int {
	typeList := make([]int, length)
	// 每种类型随机分布在列表中
	for i := 1; i <= 3; i++ {
		for {
			randIndex, _ := rand.Int(rand.Reader, big.NewInt(int64(length)))
			r := randIndex.Uint64()
			if typeList[r] == 0 {
				typeList[r] = i
				break
			}
		}
	}
	// 补全空余的列表
	for i := 0; i < len(typeList); i++ {
		randType, _ := rand.Int(rand.Reader, big.NewInt(int64(3)))
		if typeList[i] != 0 {
			continue
		}
		r := randType.Uint64()
		r++
		typeList[i] = int(r)
	}
	return typeList
}

func randomValue(t int) byte {
	switch t {
	case 1:
		b, _ := rand.Int(rand.Reader, big.NewInt(int64(len(numeric))))
		return numeric[b.Int64()]
	case 2:
		b, _ := rand.Int(rand.Reader, big.NewInt(int64(len(upper))))
		return upper[b.Int64()]
	default:
		b, _ := rand.Int(rand.Reader, big.NewInt(int64(len(lower))))
		return lower[b.Int64()]
	}
}

const (
	PKCS1PublicKeyType  = "RSA PUBLIC KEY"
	PKIXPublicKeyType   = "PUBLIC KEY"
	PKCS1PrivateKeyType = "RSA PRIVATE KEY"
)

func GenerateRSAKey(privateFile, publicFile string, bits int) error {
	// 生成 RSA 私钥对象
	privateKey, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return err
	}

	// 通过x509标准将得到的ras私钥序列化为 ASN.1 的 DER 编码字符串
	x509PrivateKey := x509.MarshalPKCS1PrivateKey(privateKey)

	// 新建私钥文件
	private, err := os.Create(privateFile)
	if err != nil {
		return err
	}
	defer private.Close()

	// 将数据保存到文件
	err = pem.Encode(private, &pem.Block{
		Type:    PKCS1PrivateKeyType,
		Headers: map[string]string{},
		Bytes:   x509PrivateKey,
	})
	if err != nil {
		return err
	}

	// X509对公钥编码
	X509PublicKey, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
	if err != nil {
		return err
	}

	// 新建公钥文件
	public, err := os.Create(publicFile)
	if err != nil {
		return err
	}
	defer public.Close()

	// 将数据保存到文件
	err = pem.Encode(public, &pem.Block{
		Type:    PKIXPublicKeyType,
		Headers: map[string]string{},
		Bytes:   X509PublicKey,
	})
	if err != nil {
		return err
	}

	return nil
}

// ReadRSAPKCS1PrivateKey 读取PKCS1格式私钥
func ReadRSAPKCS1PrivateKey(path string) (*rsa.PrivateKey, error) {
	// 读取文件
	context, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	// pem解码
	pemBlock, _ := pem.Decode(context)
	// x509解码
	privateKey, err := x509.ParsePKCS1PrivateKey(pemBlock.Bytes)
	return privateKey, err
}

// ReadRSAPublicKey 读取公钥(包含PKCS1和PKCS8)
func ReadRSAPublicKey(path string) (*rsa.PublicKey, error) {
	var err error
	// 读取文件
	readFile, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	// 使用pem解码
	pemBlock, _ := pem.Decode(readFile)
	var pkixPublicKey interface{}
	switch pemBlock.Type {
	case PKCS1PublicKeyType:
		pkixPublicKey, err = x509.ParsePKCS1PublicKey(pemBlock.Bytes)
	case PKIXPublicKeyType:
		pkixPublicKey, err = x509.ParsePKIXPublicKey(pemBlock.Bytes)
	default:
		return nil, errors.New("unknown public key type")
	}
	if err != nil {
		return nil, err
	}
	publicKey := pkixPublicKey.(*rsa.PublicKey)
	return publicKey, nil
}

// RSAEncrypt 加密(使用公钥加密)
func RSAEncrypt(data, publicKeyPath string) (string, error) {
	// 获取公钥
	// ReadRSAPublicKey代码在 【3.读取密钥】
	rsaPublicKey, err := ReadRSAPublicKey(publicKeyPath)
	if err != nil {
		return "", err
	}
	// 加密
	encryptPKCS1v15, err := rsa.EncryptPKCS1v15(rand.Reader, rsaPublicKey, []byte(data))
	if err != nil {
		return "", err
	}
	// 把加密结果转成Base64
	encryptString := base64.StdEncoding.EncodeToString(encryptPKCS1v15)
	return encryptString, err
}

// RSADecrypt 解密(使用私钥解密)
func RSADecrypt(base64data, privateKeyPath string) (string, error) {
	// data反解base64
	decodeString, err := base64.StdEncoding.DecodeString(base64data)
	if err != nil {
		return "", err
	}
	// 读取密钥
	rsaPrivateKey, err := ReadRSAPKCS1PrivateKey(privateKeyPath)
	if err != nil {
		return "", err
	}
	// 解密
	decryptPKCS1v15, err := rsa.DecryptPKCS1v15(rand.Reader, rsaPrivateKey, decodeString)
	return string(decryptPKCS1v15), err
}

func EncryptAES(key string, plainText string) (string, error) {
	hash := sha256.Sum256([]byte(key))
	aesCipher, err := aes.NewCipher(hash[:])
	if err != nil {
		return "", err
	}
	blockSize := aesCipher.BlockSize()
	pad := byte(blockSize - (len(plainText) % blockSize))
	paddedPlaintext := bytes.Repeat([]byte{pad}, int(pad))
	paddedPlaintext = append(paddedPlaintext, plainText...)
	mode := cipher.NewCBCEncrypter(aesCipher, hash[:16])
	out := make([]byte, len(paddedPlaintext))
	mode.CryptBlocks(out, paddedPlaintext)
	return hex.EncodeToString(out), nil
}

func DecryptAES(key string, encryptText string) (string, error) {
	hash := sha256.Sum256([]byte(key))
	aesCipher, err := aes.NewCipher(hash[:])
	if err != nil {
		return "", err
	}
	blockSize := aesCipher.BlockSize()
	mode := cipher.NewCBCDecrypter(aesCipher, hash[:blockSize])

	// 解密
	data, err := hex.DecodeString(encryptText)
	paddedPlaintext := make([]byte, len(data))
	mode.CryptBlocks(paddedPlaintext, data)

	// 去除填充
	pad := paddedPlaintext[0]
	plainText := paddedPlaintext[pad:]
	return string(plainText), err
}

func DecodeBase64(data string) (string, error) {
	decodeString, err := base64.StdEncoding.DecodeString(data)
	if err != nil {
		return "", err
	}
	res := string(decodeString)
	if !regexPassword.MatchString(res) {
		return "", errors.New("password contains illegal characters")
	}
	return string(decodeString), nil
}
